"""
模型发现服务
负责从订阅服务器发现新模型
"""

import asyncio
import time
from typing import List, Dict, Any, Optional, Set
import aiohttp

from ..storage import data_manager, SubscriptionServer, ModelInfo
from ..logger import logger
from ..config import config
from ..utils import retry_on_exception, generate_id, compare_versions


class ModelDiscoveryService:
    """模型发现服务"""
    
    def __init__(self):
        self.subscription_config = config.get_subscription_config()
        self.is_running = False
        self.discovery_task = None
        self.discovered_models: Dict[str, Set[str]] = {}  # server_id -> set of model_ids
        
    async def start(self):
        """启动发现服务"""
        if self.is_running:
            return
            
        self.is_running = True
        
        if self.subscription_config.get('enabled', True):
            # 启动定期发现任务
            self.discovery_task = asyncio.create_task(self._periodic_discovery())
            logger.info("Model discovery service started")
        else:
            logger.info("Model discovery service is disabled")
            
    async def stop(self):
        """停止发现服务"""
        self.is_running = False
        
        if self.discovery_task:
            self.discovery_task.cancel()
            try:
                await self.discovery_task
            except asyncio.CancelledError:
                pass
                
        logger.info("Model discovery service stopped")
        
    async def _periodic_discovery(self):
        """定期发现模型"""
        default_interval = self.subscription_config.get('default_check_interval', 1800)
        
        while self.is_running:
            try:
                await self.discover_all_servers()
                await asyncio.sleep(default_interval)
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Error in periodic model discovery: {e}")
                await asyncio.sleep(60)  # 出错时等待1分钟再重试
                
    async def discover_all_servers(self):
        """发现所有服务器的模型"""
        try:
            # 获取所有在线服务器
            servers = await data_manager.get_servers()
            online_servers = [server for server in servers if server.status == "online"]
            
            if not online_servers:
                logger.debug("No online servers for model discovery")
                return
                
            logger.info(f"Discovering models from {len(online_servers)} servers")
            
            # 并发发现各个服务器的模型
            max_concurrent = self.subscription_config.get('max_concurrent_checks', 5)
            semaphore = asyncio.Semaphore(max_concurrent)
            
            tasks = []
            for server in online_servers:
                task = asyncio.create_task(
                    self._discover_server_models(semaphore, server)
                )
                tasks.append(task)
            
            # 等待所有发现完成
            results = await asyncio.gather(*tasks, return_exceptions=True)
            
            # 统计结果
            successful_discoveries = sum(1 for result in results if not isinstance(result, Exception))
            logger.info(f"Model discovery completed: {successful_discoveries}/{len(online_servers)} servers")
            
        except Exception as e:
            logger.error(f"Failed to discover models: {e}")
            
    async def _discover_server_models(self, semaphore: asyncio.Semaphore, server: SubscriptionServer):
        """发现特定服务器的模型"""
        async with semaphore:
            try:
                logger.debug(f"Discovering models from server: {server.name}")
                
                # 获取服务器模型列表
                server_models = await self._fetch_server_models(server)
                if server_models is None:
                    logger.warning(f"Failed to fetch models from server {server.name}")
                    return
                
                # 处理发现的模型
                new_models_count = await self._process_discovered_models(server, server_models)
                
                if new_models_count > 0:
                    logger.info(f"Discovered {new_models_count} new models from {server.name}")
                else:
                    logger.debug(f"No new models found on {server.name}")
                    
            except Exception as e:
                logger.error(f"Error discovering models from server {server.name}: {e}")
                
    @retry_on_exception(max_retries=3, delay=1.0)
    async def _fetch_server_models(self, server: SubscriptionServer) -> Optional[List[Dict[str, Any]]]:
        """从服务器获取模型列表"""
        try:
            timeout = aiohttp.ClientTimeout(
                total=self.subscription_config.get('request_timeout', 30)
            )
            
            headers = {}
            if server.api_key:
                headers['Authorization'] = f'Bearer {server.api_key}'
            
            async with aiohttp.ClientSession(timeout=timeout) as session:
                models_url = f"{server.url.rstrip('/')}/api/v1/models"
                
                async with session.get(models_url, headers=headers) as response:
                    if response.status == 200:
                        data = await response.json()
                        return data.get('models', [])
                    else:
                        logger.warning(f"Server {server.name} returned status {response.status}")
                        return None
                        
        except Exception as e:
            logger.debug(f"Failed to fetch models from {server.name}: {e}")
            raise
            
    async def _process_discovered_models(
        self, 
        server: SubscriptionServer, 
        server_models: List[Dict[str, Any]]
    ) -> int:
        """处理发现的模型"""
        try:
            new_models_count = 0
            
            # 获取当前已知的模型
            existing_models = await data_manager.get_models()
            existing_model_keys = set()
            
            for model in existing_models:
                if model.server_id == server.id:
                    key = f"{model.name}:{model.version}"
                    existing_model_keys.add(key)
            
            # 处理每个发现的模型
            for model_data in server_models:
                try:
                    model_name = model_data.get('name', '')
                    model_version = model_data.get('version', '1.0.0')
                    model_key = f"{model_name}:{model_version}"
                    
                    # 检查是否是新模型
                    if model_key not in existing_model_keys:
                        # 创建新模型记录
                        new_model = await self._create_model_from_data(server, model_data)
                        if new_model:
                            await data_manager.add_model(new_model)
                            new_models_count += 1
                            
                            # 触发模型发现事件
                            await self._trigger_model_discovered_event(server, new_model)
                            
                except Exception as e:
                    logger.error(f"Error processing model {model_data}: {e}")
                    
            return new_models_count
            
        except Exception as e:
            logger.error(f"Error processing discovered models: {e}")
            return 0
            
    async def _create_model_from_data(
        self, 
        server: SubscriptionServer, 
        model_data: Dict[str, Any]
    ) -> Optional[ModelInfo]:
        """从服务器数据创建模型信息"""
        try:
            current_time = time.time()
            
            model = ModelInfo(
                id=generate_id(),
                name=model_data.get('name', ''),
                version=model_data.get('version', '1.0.0'),
                model_type=model_data.get('model_type', 'detection'),
                description=model_data.get('description', ''),
                author=model_data.get('author', ''),
                size=model_data.get('size', 'm'),
                upload_time=model_data.get('upload_time', current_time),
                last_update=current_time,
                status=model_data.get('status', 'available'),
                download_url=model_data.get('download_url', ''),
                checksum=model_data.get('checksum', ''),
                server_id=server.id,
                tags=model_data.get('tags', []),
                requirements=model_data.get('requirements', {}),
                metrics=model_data.get('metrics', {})
            )
            
            return model
            
        except Exception as e:
            logger.error(f"Error creating model from data: {e}")
            return None
            
    async def _trigger_model_discovered_event(self, server: SubscriptionServer, model: ModelInfo):
        """触发模型发现事件"""
        try:
            # 通知订阅服务有新模型发现
            from .notification_service import notification_service
            
            event_data = {
                "event_type": "model_discovered",
                "server_id": server.id,
                "server_name": server.name,
                "model_id": model.id,
                "model_name": model.name,
                "model_version": model.version,
                "model_type": model.model_type,
                "timestamp": time.time()
            }
            
            await notification_service.broadcast_event(event_data)
            
            logger.info(f"Model discovered event triggered: {model.name} v{model.version} from {server.name}")
            
        except Exception as e:
            logger.error(f"Error triggering model discovered event: {e}")
            
    async def discover_server_now(self, server_id: str) -> Dict[str, Any]:
        """立即发现指定服务器的模型"""
        try:
            # 获取服务器信息
            server = await data_manager.get_server_by_id(server_id)
            if not server:
                raise ValueError(f"Server {server_id} not found")
            
            # 发现模型
            server_models = await self._fetch_server_models(server)
            if server_models is None:
                return {
                    "success": False,
                    "message": f"Failed to fetch models from server {server.name}",
                    "new_models": 0
                }
            
            # 处理发现的模型
            new_models_count = await self._process_discovered_models(server, server_models)
            
            return {
                "success": True,
                "message": f"Discovery completed for server {server.name}",
                "total_models": len(server_models),
                "new_models": new_models_count
            }
            
        except Exception as e:
            logger.error(f"Error in immediate discovery for server {server_id}: {e}")
            return {
                "success": False,
                "message": str(e),
                "new_models": 0
            }


# 全局模型发现服务实例
discovery_service = ModelDiscoveryService()
